import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mplfinance.original_flavor import candlestick_ohlc


def calculate_kangaroo_tail(data, window=200):
    """
    判断K线图中的袋鼠尾形态

    参数:
    data: pandas DataFrame, 包含开盘价(open), 最高价(high), 最低价(low), 收盘价(close)
    window: int, 计算平均实体长度的窗口大小, 默认为200

    返回:
    DataFrame: 添加了袋鼠尾标记的原始数据
    """
    # 计算每根K线的实体长度(绝对值)
    data['body_length'] = np.abs(data['close'] - data['open'])

    # 计算每根K线的总长度(影线+实体)
    data['total_length'] = data['high'] - data['low']

    # 计算实体占总长度的比例
    data['body_ratio'] = data['body_length'] / data['total_length'].replace(0, 0.0001)  # 避免除以0

    # 计算近window根K线的平均实体长度
    data['avg_body_length'] = data['body_length'].rolling(window=window, min_periods=1).mean()

    # 初始化袋鼠尾标记列
    data['kangaroo_tail'] = 0  # 0表示不是袋鼠尾, 1表示上升袋鼠尾, -1表示下降袋鼠尾

    for i in range(1, len(data) - 1):
        # 检查当前K线是否符合袋鼠尾的基本条件
        current_avg = data['avg_body_length'].iloc[i]

        # 实体比例小于20%且总长度是平均的3倍以上
        if (data['body_ratio'].iloc[i] < 0.2 and
                data['total_length'].iloc[i] > 3 * current_avg):

            # 检查第二天是否符合条件(高度是正常高度的80%-120%)
            next_length = data['total_length'].iloc[i + 1]
            if 0.8 * current_avg <= next_length <= 1.2 * current_avg:

                # 判断是上升袋鼠尾还是下降袋鼠尾
                # 上升袋鼠尾: 实体在上部(下影线长)
                if (data['close'].iloc[i] - data['open'].iloc[i]) > 0:
                    # 检查开盘价是否略低于前收盘价
                    prev_close = data['close'].iloc[i - 1]
                    current_open = data['open'].iloc[i]
                    if current_open < prev_close and (prev_close - current_open) <= 0.01 * prev_close:
                        data.at[data.index[i], 'kangaroo_tail'] = 1

                # 下降袋鼠尾: 实体在下部(上影线长)
                else:
                    # 检查开盘价是否略高于前收盘价
                    prev_close = data['close'].iloc[i - 1]
                    current_open = data['open'].iloc[i]
                    if current_open > prev_close and (current_open - prev_close) <= 0.01 * prev_close:
                        data.at[data.index[i], 'kangaroo_tail'] = -1

    return data


# 示例使用
if __name__ == "__main__":
    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

    # 假设我们有一个包含OHLC数据的DataFrame
    # 这里生成一些随机数据作为示例
    np.random.seed(42)
    dates = pd.date_range(start='2020-01-01', periods=300)
    opens = np.random.normal(100, 2, 300).cumsum()
    highs = opens + np.random.uniform(0.5, 3, 300)
    lows = opens - np.random.uniform(0.5, 3, 300)
    closes = opens + np.random.normal(0, 1, 300)

    # 创建一个明显的上升袋鼠尾形态
    closes[150] = opens[150] + 0.05  # 非常小的实体
    highs[150] = opens[150] + 20  # 很长的上影线
    lows[150] = opens[150] - 0.5  # 短下影线

    # 第二天
    opens[151] = closes[150] - 0.1  # 略低于前收盘价
    closes[151] = opens[151] + 2  # 正常高度的实体
    highs[151] = opens[151] + 2.5  # 正常范围的高度
    lows[151] = opens[151] - 0.5  # 正常范围的低点

    data = pd.DataFrame({
        'date': dates,
        'open': opens,
        'high': highs,
        'low': lows,
        'close': closes
    }).set_index('date')

    # 计算袋鼠尾
    result = calculate_kangaroo_tail(data)

    # 查看结果
    print("检测到的袋鼠尾形态:")
    print(result[result['kangaroo_tail'] != 0])

    # 可视化
    fig, ax = plt.subplots(figsize=(12, 6))
    subset = result.iloc[140:160].reset_index()
    subset['date_num'] = subset.index

    candlestick_ohlc(ax, subset[['date_num', 'open', 'high', 'low', 'close']].values,
                     width=0.6, colorup='g', colordown='r')

    # 标记袋鼠尾
    kangaroo_idx = subset[subset['kangaroo_tail'] != 0].index
    for idx in kangaroo_idx:
        ax.annotate('袋鼠尾', xy=(idx, subset['high'].iloc[idx]),
                    xytext=(0, 20), textcoords='offset points',
                    arrowprops=dict(arrowstyle='->'), ha='center')

    plt.title('K线图与袋鼠尾识别')
    plt.xlabel('时间')
    plt.ylabel('价格')
    plt.grid()
    plt.show()